# -*- coding: UTF-8 -*-
"""
@Project ：swift-forge-engine 
@File    ：generator
@IDE     ：PyCharm 
@Author  ：Tony
@Date    ：2025/2/19 16:15 
@Des     ：
"""
import re
from pathlib import Path
from jinja2 import Environment, FileSystemLoader
from sqlalchemy import inspect, create_engine, URL, MetaData


class TemplateBasedGenerator:
    def __init__(self, db_config, template_dir, output_dir, base_package, include_tables=None):
        self.db_config = db_config
        self.engine = create_engine(URL.create(**db_config))
        self.metadata = MetaData()
        self.metadata.reflect(self.engine)

        # 表过滤
        self.include_tables = set(include_tables) if include_tables else None

        # 配置Jinja2环境
        self.env = Environment(
            loader=FileSystemLoader(template_dir),
            trim_blocks=True,
            lstrip_blocks=True,
            keep_trailing_newline=True
        )
        # 基础包路径参数
        self.base_package = base_package
        # 创建多模块目录
        self.models_dir = Path(output_dir) / "models"
        self.schemas_dir = Path(output_dir) / "schemas"
        self.crud_dir = Path(output_dir) / "crud"
        self.services_dir = Path(output_dir) / "services"
        self.api_dir = Path(output_dir) / "apis/v1/sys"

        # 确保目录存在
        for d in [self.models_dir, self.schemas_dir, self.crud_dir, self.services_dir, self.api_dir]:
            d.mkdir(parents=True, exist_ok=True)
            init_file = d / "__init__.py"
            if not init_file.exists():
                init_file.touch()

    def _get_import_path(self, module_type: str, table_name: str) -> str:
        """生成动态导入路径"""
        paths = {
            "crud": f"{self.base_package}.{self.crud_dir.name}.{table_name}_crud",
            "model": f"{self.base_package}.{self.models_dir.name}.{table_name}",
            "schema": f"{self.base_package}.{self.schemas_dir.name}.{table_name}_schema",
            "api": f"{self.base_package}.{self.api_dir.name}.{table_name}",
            "services": f"{self.base_package}.{self.services_dir.name}.{table_name}_service"
        }
        return paths[module_type]

    def generate_all(self):
        """生成所有模块"""
        self.generate_models()
        self.generate_schemas()
        self.generate_crud()
        self.generate_services()
        self.generate_apis()

    def generate_models(self):
        """表的模型类"""
        inspector = inspect(self.engine)
        template = self.env.get_template('model.jinja2')

        # 获取过滤后的表列表
        tables = self._get_filtered_tables()

        for table_name in tables:
            # 准备模板上下文数据
            context = self._prepare_context(inspector, table_name)

            # 渲染模板
            code = template.render(context)

            # 写入文件
            output_path = self.models_dir / f"{table_name}.py"
            output_path.write_text(code, encoding='utf-8')

    def _prepare_context(self, inspector, table_name):
        """准备模板需要的上下文数据"""
        return {
            'database_type': self.db_config['drivername'].split('+')[0],
            'default_datetime_column': True,  # 根据实际情况调整
            'table_name_class': self._to_camel_case(table_name),
            'table_name_zh': self._get_table_comment(inspector, table_name),
            'table_name_en': table_name,
            'models': self._get_column_definitions(inspector, table_name)
        }

    def generate_schemas(self):
        """生成Schema类（与您原有代码兼容的改进版）"""
        inspector = inspect(self.engine)
        template = self.env.get_template('schema.jinja2')

        for table_name in self._get_filtered_tables():
            context = self._prepare_schema_context(inspector, table_name)
            code = template.render(context)

            output_path = self.schemas_dir / f"{table_name}_schema.py"
            output_path.write_text(code, encoding='utf-8')

    def _prepare_schema_context(self, inspector, table_name):
        """准备Schema生成专用上下文"""
        base_context = self._prepare_context(inspector, table_name)
        columns = inspector.get_columns(table_name)
        pk_info = inspector.get_pk_constraint(table_name)

        # 添加Schema专用字段
        base_context.update({
            'primary_keys': pk_info['constrained_columns'],
            'has_created_time': any(col['name'] == 'created_time' for col in columns),
            'has_updated_time': any(col['name'] == 'updated_time' for col in columns),
            'column_type_map': {col['name']: self._map_python_type(str(col['type'])) for col in columns}
        })
        return base_context

    def generate_crud(self):
        """生成CRUD操作类"""
        inspector = inspect(self.engine)
        template = self.env.get_template('crud.jinja2')

        tables = self._get_filtered_tables()
        for table_name in tables:
            context = self._prepare_crud_context(inspector, table_name)
            code = template.render(context)

            output_path = self.crud_dir / f"{table_name}_crud.py"
            output_path.write_text(code, encoding='utf-8')

    def _prepare_crud_context(self, inspector, table_name):
        """准备CRUD生成专用上下文"""
        base_context = self._prepare_context(inspector, table_name)

        return {
            **base_context,
            "model_import_path": self._get_import_path("model", table_name),
            "schema_import_path": self._get_import_path("schema", table_name),
            'table_simple_name_zh': self._get_simple_chinese_name(base_context['table_name_zh']),
            "schema_name": base_context["table_name_class"]
        }

    def generate_services(self):
        """生成服务层代码"""
        inspector = inspect(self.engine)
        template = self.env.get_template('services.jinja2')

        for table_name in self._get_filtered_tables():
            context = self._prepare_service_context(inspector, table_name)
            code = template.render(context)

            output_path = self.services_dir / f"{table_name}_service.py"
            output_path.write_text(code, encoding='utf-8')

    def _prepare_service_context(self, inspector, table_name):
        """准备服务层上下文"""
        base_context = self._prepare_context(inspector, table_name)

        return {
            **base_context,
            'crud_import_path': self._get_import_path('crud', table_name),
            'model_import_path': self._get_import_path('model', table_name),
            'schema_import_path': self._get_import_path('schema', table_name),
            'primary_key_type': self._get_primary_key_type(inspector, table_name),
            'table_simple_name_zh': self._get_simple_chinese_name(base_context['table_name_zh']),
            "schema_name": base_context["table_name_class"]
        }

    def generate_apis(self):
        """生成API路由文件"""
        inspector = inspect(self.engine)
        template = self.env.get_template('api.jinja2')

        for table_name in self._get_filtered_tables():
            context = self._prepare_api_context(inspector, table_name)
            code = template.render(context)

            output_path = self.api_dir / f"{table_name}.py"
            output_path.write_text(code, encoding='utf-8')

    def _prepare_api_context(self, inspector, table_name):
        """准备API生成上下文"""
        base_context = self._prepare_context(inspector, table_name)
        pk_info = inspector.get_pk_constraint(table_name)

        return {
            **base_context,
            'permission_prefix': self._get_permission_prefix(table_name),
            'primary_key': self._get_primary_key_name(pk_info),
            'primary_key_type': self._get_primary_key_type(inspector, table_name),
            'service_import_path': self._get_import_path('services', table_name),
            'schema_import_path': self._get_import_path('schema', table_name),
            "schema_name": base_context["table_name_class"]
        }

    def _get_permission_prefix(self, table_name):
        """转换表名为权限前缀"""
        return table_name.replace('_', ':').replace('sys:', '')

    def _get_primary_key_name(self, pk_info):
        """获取主键字段名"""
        return pk_info['constrained_columns'][0] if pk_info['constrained_columns'] else 'id'

    def _get_primary_key_type(self, inspector, table_name):
        """获取主键类型"""
        for col in inspector.get_columns(table_name):
            if col['name'] in inspector.get_pk_constraint(table_name)['constrained_columns']:
                return self._map_python_type(str(col['type']).split('(')[0])
        return 'int'

    def _get_simple_chinese_name(self, raw_comment):
        """提取简化的中文表名"""
        if not raw_comment:
            return "数据"
        # 去除注释中的特殊字符
        return re.sub(r'[表\s]', '', raw_comment.split()[-1]) or "数据"

    def _get_filtered_tables(self):
        """获取需要处理的表列表"""
        all_tables = list(self.metadata.tables.keys())

        # 未指定表时返回全部
        if self.include_tables is None:
            return all_tables

        # 过滤并验证表存在性
        valid_tables = []
        for table in self.include_tables:
            if table in all_tables:
                valid_tables.append(table)
            else:
                print(f"警告：表 {table} 不存在，已跳过")

        return valid_tables

    def _get_column_definitions(self, inspector, table_name):
        """转换列定义为模板需要的格式"""
        columns = []
        for idx, column in enumerate(inspector.get_columns(table_name)):
            col_type = str(column['type'])
            is_nullable = column['nullable']

            # 解析字段类型和长度
            sa_type, length = self._parse_column_type(col_type)

            # 转换默认值
            default_value = self._convert_default(column.get('default'))

            # 构建模型定义
            column_def = {
                'name': column['name'],
                'type': sa_type,
                'pd_type': self._map_python_type(sa_type),
                'length': length,
                'is_nullable': is_nullable,
                'default': default_value,
                'comment': column.get('comment'),
                'sort': idx + 1
            }
            columns.append(column_def)
        return columns

    def _parse_column_type(self, column_type):
        """解析SQLAlchemy类型为基本类型和长度"""
        type_str = str(column_type).split('(')[0]
        length = None

        # 提取长度信息
        if '(' in str(column_type):
            length = int(str(column_type).split('(')[1].split(')')[0])

        # 类型映射
        type_map = {
            'INT': 'Integer',
            'VARCHAR': 'String',
            'TEXT': 'Text',
            'DATETIME': 'DateTime',
            'VARBINARY': 'VARBINARY',
            'TINYINT': 'Boolean'
        }
        return type_map.get(type_str.upper(), type_str), length

    def _map_python_type(self, sa_type):
        """映射SQLAlchemy类型到Python类型"""
        type_mapping = {
            'Integer': 'int',
            'String': 'str',
            'Text': 'str',
            'DateTime': 'datetime',
            'Date': 'date',
            'Boolean': 'bool',
            'Float': 'float',
            'VARBINARY': 'bytes'
        }
        return type_mapping.get(sa_type, 'str')

    def _convert_default(self, default_value):
        """转换数据库默认值为Python可用的值"""
        if not default_value:
            return None

        # 处理函数默认值
        if 'nextval' in default_value:
            return None  # 忽略序列
        if 'CURRENT_TIMESTAMP' in default_value:
            return 'timezone.now()'
        if 'uuid_generate' in default_value:
            return 'uuid4_str()'

        # 去除字符串引号
        if isinstance(default_value, str) and default_value.startswith("'"):
            return default_value.strip("'")

        return default_value

    def _to_camel_case(self, name):
        """转换为驼峰命名"""
        return ''.join(word.capitalize() for word in name.split('_'))

    def _get_table_comment(self, inspector, table_name):
        """获取表注释"""
        try:
            # 获取原始注释
            raw_comment = inspector.get_table_comment(table_name)

            # 处理不同数据库返回格式
            if isinstance(raw_comment, dict):
                # MySQL返回格式 {'text': 'comment'}
                return raw_comment.get('text', table_name)
            elif isinstance(raw_comment, str):
                # PostgreSQL直接返回字符串
                return raw_comment.strip() or table_name
            else:
                return table_name
        except Exception as e:
            print(f"获取表注释错误 {table_name}: {str(e)}")
            return table_name


if __name__ == '__main__':
    # 使用示例
    generator = TemplateBasedGenerator(
        db_config={
            'drivername': 'mysql+pymysql',
            'host': 'rm-uf68nrr3cqqu2vi4eio.mysql.rds.aliyuncs.com',
            'username': 'tony',
            'password': 'Hanxun_5200',
            'database': 'sfe_test'
        },
        template_dir='templates',
        base_package="App.app.admin",
        output_dir='App/app/admin',
        include_tables=['sys_config']
    )
    generator.generate_all()
